{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6321595e",
   "metadata": {},
   "source": [
    "# Cox生存分析\n",
    "\n",
    "* `mydir`：自己的数据\n",
    "* `ostime_column`: 数据对应的生存时间，不一定非的是OST，也可以是DST、FST等。\n",
    "* `os`：生存状态，不一定非的是OS，也可以是DS、FS等。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb0fc498",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lifelines import CoxPHFitter\n",
    "import pandas as pd\n",
    "from onekey_algo.custom.components.comp1 import normalize_df\n",
    "from sklearn.model_selection import train_test_split\n",
    "from onekey_algo import get_param_in_cwd\n",
    "from onekey_algo.custom.components.comp1 import fillna\n",
    "\n",
    "event_col = get_param_in_cwd('event_col')\n",
    "group_info = 'group'\n",
    "task_type = 'Clinical_'\n",
    "duration_col= get_param_in_cwd('duration_col')\n",
    "data = pd.read_csv('data/clinical.csv', dtype={'ID': str})\n",
    "data = data[data.columns[:-3]]\n",
    "# data = normalize_df(data, not_norm='ID')\n",
    "label_data = pd.read_csv(get_param_in_cwd('label_file'), dtype={'ID': str})\n",
    "label_data['ID'] = label_data['ID'].map(lambda x: f\"{x}.nii.gz\" if not (f\"{x}\".endswith('.nii.gz') or  f\"{x}\".endswith('.nii')) else x)\n",
    "data = pd.merge(data, label_data[['ID', event_col, duration_col, 'group']], on='ID', how='inner')\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f54b65a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import onekey_algo.custom.components as okcomp\n",
    "from collections import OrderedDict\n",
    "\n",
    "\n",
    "train_data = data[(data[group_info] == 'train')]\n",
    "\n",
    "# subsets = [s for s in label_data['group'].value_counts().index if s != 'train']\n",
    "subsets = get_param_in_cwd('subsets')\n",
    "val_datasets = OrderedDict()\n",
    "for subset in subsets:\n",
    "    val_data = data[data[group_info] == subset]\n",
    "    val_datasets[subset] = val_data\n",
    "    val_data.to_csv(f'features/{task_type}{subset}_features_norm.csv', index=False)\n",
    "\n",
    "print('，'.join([f\"{subset}样本数：{d_.shape}\" for subset, d_ in val_datasets.items()]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b265817c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from onekey_algo.custom.components.survival import uni_cox\n",
    "\n",
    "if os.path.exists(f'features/{task_type}features_unisel.csv') and False:\n",
    "    train_data = pd.read_csv(f'features/{task_type}features_unisel.csv')\n",
    "else:\n",
    "    sel_features, uni_info = uni_cox(train_data, duration_col=duration_col, event_col=event_col,\n",
    "                                     cols=[c for c in train_data.columns if c not in [event_col, duration_col, 'ID', 'group']], \n",
    "                                     pvalue_thres=0.05, verbose=True)\n",
    "    train_data = train_data[['ID'] + sel_features + [event_col, duration_col, 'group']]\n",
    "    data[['ID'] + sel_features + [event_col, duration_col, 'group']].to_csv(f'features/{task_type}features_unisel.csv', header=True, index=False)\n",
    "train_data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cce536d",
   "metadata": {},
   "source": [
    "## Cox概览\n",
    "\n",
    "所有Cox回归的必要数据，主要关注的数据有3个\n",
    "1. `Concordance`: c-index\n",
    "2. `exp(coef)`: 每个特征对应的HR，同时也有期对应的95%分位数。\n",
    "3. `p`: 表示特征是否显著。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0201324",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from lifelines import CoxPHFitter\n",
    "\n",
    "cph = CoxPHFitter(penalizer=0.5)\n",
    "cph.fit(train_data[[c for c in train_data.columns if c not in ['ID', 'group']]], duration_col=duration_col, event_col=event_col)\n",
    "cph.print_summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "baddb16e",
   "metadata": {},
   "outputs": [],
   "source": [
    "multi_info = cph.summary[['exp(coef)', 'exp(coef) lower 95%', 'exp(coef) upper 95%', 'p']].reset_index()\n",
    "i = pd.merge(uni_info, multi_info, on='covariate', how='left', suffixes=['_Uni', '_Multi'])\n",
    "multi_data = data[['ID'] + list(multi_info[multi_info['p'] < 0.05]['covariate'])]\n",
    "multi_data.to_csv(f'features/{task_type}features_mulsel.csv', header=True, index=False)\n",
    "i['p_Uni'] = i['p_Uni'].map(lambda x: '<0.05' if not pd.isna(x) and x < 0.05 else f\"{x:.3f}\")\n",
    "i['p_Multi'] = i['p_Multi'].map(lambda x: '<0.05' if not pd.isna(x) and x < 0.05 else f\"{x:.3f}\")\n",
    "i.columns = ['features_name', 'HR', 'lower 95%CI', 'upper 95%CI', 'pvalue', 'HR', 'lower 95%CI', 'upper 95%CI', 'pvalue']\n",
    "i = i.applymap(lambda x: '' if pd.isna(x) else x)\n",
    "i.to_csv('results/unimulti.csv', index=False)\n",
    "i"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac98781f",
   "metadata": {},
   "source": [
    "#### 输出每个特征的HR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "349aa141",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.figure(figsize=(10, train_data.shape[1]-4))\n",
    "cph.plot(hazard_ratios=True)\n",
    "plt.savefig('img/Clinic_feature_pvalue.svg')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bf8910e",
   "metadata": {},
   "source": [
    "# KM 曲线\n",
    "\n",
    "根据HR进行分组，计算KM以及log ranktest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "669023fa",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from lifelines import CoxPHFitter\n",
    "from lifelines.statistics import logrank_test\n",
    "from lifelines import KaplanMeierFitter\n",
    "from lifelines.plotting import add_at_risk_counts\n",
    "\n",
    "thres = 1e-4\n",
    "bst_split = {'train': 1.0, 'val':1, 'test': 1}\n",
    "loc = {'train': 0.5, 'val':0.5, 'test': 0.5}\n",
    "\n",
    "for subset, test_data in val_datasets.items():\n",
    "    c_index = cph.score(test_data[[c for c in test_data.columns if c != 'ID']], scoring_method=\"concordance_index\")\n",
    "#     y_pred = cph.predict_median(test_data[[c for c in test_data.columns if c != 'ID']])\n",
    "#     cox_data = pd.concat([test_data, y_pred], axis=1)\n",
    "#     mean = cox_data.describe()[0.5]['mean']\n",
    "#     cox_data['HR'] = cox_data[0.5] < mean\n",
    "    y_pred = cph.predict_partial_hazard(test_data[[c for c in test_data.columns if c != 'ID']])\n",
    "    cox_data = pd.concat([test_data, y_pred], axis=1)\n",
    "    mean = cox_data.describe()[0]['50%']\n",
    "    cox_data['HR'] = cox_data[0] > mean\n",
    "#     cox_data['HR'] = cox_data[0] > 1\n",
    "\n",
    "    dem = (cox_data[\"HR\"] == True)\n",
    "    results = logrank_test(cox_data[duration_col][dem], cox_data[duration_col][~dem], \n",
    "                           event_observed_A=cox_data[event_col][dem], event_observed_B=cox_data[event_col][~dem])\n",
    "    p_value = f\"={results.p_value:.3f}\" if results.p_value > thres else f'<{thres}'\n",
    "    plt.title(f\"Cohort {subset} C-index:{c_index:.3f}\")\n",
    "    plt.ylabel('Probability')\n",
    "    if sum(dem):\n",
    "        kmf_high = KaplanMeierFitter()\n",
    "        kmf_high.fit(cox_data[duration_col][dem], event_observed=cox_data[event_col][dem], label=\"High Risk\")\n",
    "        kmf_high.plot_survival_function(color='r')\n",
    "    if sum(~dem):\n",
    "        kmf_low = KaplanMeierFitter()\n",
    "        kmf_low.fit(cox_data[duration_col][~dem], event_observed=cox_data[event_col][~dem], label=\"Low Risk\")\n",
    "        kmf_low.plot_survival_function(color='g')\n",
    "    plt.text(0.5, loc[subset] if subset in loc else 0.2, f\"P{p_value}\")\n",
    "    plt.xlabel('Time(months)')\n",
    "    add_at_risk_counts(kmf_high, kmf_low, rows_to_show=['At risk'])\n",
    "    plt.savefig(f'img/{task_type}KM_{subset}.svg', bbox_inches='tight')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45e697b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "\n",
    "def get_prediction(model: CoxPHFitter, data, ID=None, **kwargs):\n",
    "    hr = model.predict_partial_hazard(data)\n",
    "    expectation = model.predict_expectation(data)\n",
    "    \n",
    "    predictions = pd.concat([hr, expectation], axis=1)\n",
    "    predictions.columns = ['HR', 'expectation']\n",
    "    if ID is not None:\n",
    "        predictions = pd.concat([ID, hr, expectation], axis=1)\n",
    "        predictions.columns = ['ID', 'HR', 'expectation']\n",
    "    else:\n",
    "        predictions = pd.concat([hr, expectation], axis=1)\n",
    "        predictions.columns = ['HR', 'expectation']\n",
    "    return predictions\n",
    "\n",
    "os.makedirs('results', exist_ok=True)\n",
    "info = []\n",
    "for subset, test_data in val_datasets.items():\n",
    "    if subset in get_param_in_cwd('subsets'):\n",
    "        results = get_prediction(cph, test_data, ID=test_data['ID'])\n",
    "        results.to_csv(f'results/{task_type}cox_predictions_{subset}.csv', index=False)\n",
    "        results['group'] = subset\n",
    "        info.append(results)\n",
    "        pd.merge(results, label_data[['ID', event_col, duration_col]], on='ID', how='inner').to_csv(f'features/{task_type}4xtile_{subset}.txt', \n",
    "                                                                                                    index=False, sep='\\t')\n",
    "info = pd.concat(info, axis=0)\n",
    "info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82433b3b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
